import os
import hashlib
import sqlite3
from datetime import datetime
import argparse
from pathlib import Path  # 导入Path类
from tqdm import tqdm

def is_video_file(filename):
    """判断给定的文件是否为视频文件"""
    video_extensions = {'.mp4', '.avi', '.mkv', '.mov', '.wmv', '.flv','.mts','.mpeg','.rm', '.webm','.3gp','.mp3','.wav','.f4v','.m4a','.mpg','.rmvb','.vob'}
    return Path(filename).suffix.lower() in video_extensions

def calculate_file_hash(file_path):
    """计算文件的SHA256哈希值"""
    hasher = hashlib.sha256()
    with open(file_path, 'rb') as f:
        buf = f.read()
        hasher.update(buf)
    return hasher.hexdigest()

def check_and_insert(db_conn, file_path, video_tmp_dir):
    """检查hash是否存在，若不存在则插入新记录"""
    cursor = db_conn.cursor()
    file_hash = calculate_file_hash(file_path)

    # 首先检查video_hash表
    cursor.execute('SELECT 1 FROM video_hash WHERE file_hash=?', (file_hash,))
    if cursor.fetchone():
        return False, None, None  # 如果在video_hash中找到，直接返回不插入
    
    # 然后检查file_merge表
    cursor.execute('SELECT 1 FROM file_merge WHERE file_hash=?', (file_hash,))
    if cursor.fetchone():
        return False, None, None  # 如果在file_merge中找到，也直接返回不插入
    
    # 计算aim_folder
    mod_time = os.path.getmtime(file_path)
    mod_time_str = datetime.fromtimestamp(mod_time).strftime('%Y')
    aim_folder = f'video{mod_time_str}/others'
    
    absolute_file_path = os.path.join(video_tmp_dir, os.path.relpath(file_path, video_tmp_dir))
    file_name = os.path.basename(file_path)
    now = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
    cursor.execute('''
        INSERT INTO file_merge (file_path, file_name, file_hash, aim_folder, deal_flag, create_time) 
        VALUES (?, ?, ?, ?, ?, ?)
    ''', (absolute_file_path, file_name, file_hash, aim_folder, 0, now))
    db_conn.commit()
    print(f"已添加文件: {file_path}")
    return True, file_name, file_hash

def scan_directory_and_merge(directory, db_conn):
    """扫描目录并处理文件"""

    # 计算总文件数，用于进度条
    total_files = sum(len(files) for _, _, files in os.walk(directory))
    
    with tqdm(total=total_files, desc='Processing files', unit='file') as pbar:
        for root, dirs, files in os.walk(directory):
            for name in files:
                file_path = os.path.join(root, name)
                # 检查是否为视频文件
                if is_video_file(file_path):
                    inserted, file_name, file_hash = check_and_insert(db_conn, file_path, directory)
                    
                pbar.update(1)

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="Scan directory and store video file information into the database.")
    parser.add_argument('--db', required=True, help="Directory containing the SQLite database file (file_hash.db)")
    parser.add_argument('--scan-dir', required=True, help="Directory to scan for video files")

    args = parser.parse_args()

    script_dir = args.db
    video_tmp_dir = args.scan_dir

    # 打印 script_dir 和 video_tmp_dir
    print(f"Script Directory: {script_dir}")
    print(f"Video Temp Directory: {video_tmp_dir}")

    # 假设数据库文件位于脚本所在目录下的'E_RelaxArea_video'文件夹中
    db_path = os.path.join(script_dir, 'file_hash.db')

    # 创建/连接数据库
    conn = sqlite3.connect(db_path)

    # 确保所有必要的数据库表存在
    cursor = conn.cursor()
    cursor.execute('''
        CREATE TABLE IF NOT EXISTS file_merge (
            id INTEGER PRIMARY KEY,
            file_path TEXT NOT NULL,
            file_name TEXT NOT NULL,
            file_hash TEXT NOT NULL,
            aim_folder TEXT NOT NULL,
            deal_flag INTEGER NOT NULL,
            create_time TEXT NOT NULL
        )
    ''')
    cursor.execute('''
        CREATE TABLE IF NOT EXISTS video_hash (
            file_hash TEXT PRIMARY KEY
        )
    ''')
    conn.commit()

    # 扫描目录并合并记录
    scan_directory_and_merge(video_tmp_dir, conn)

    # 关闭数据库连接
    conn.close()